import os.path
import sys
from PyQt5.QtWidgets import (
    QApplication, QWidget, QLabel, QLineEdit, QPushButton, QTextEdit,
    QVBoxLayout, QHBoxLayout, QSpinBox, QFileDialog, QScrollArea, QFrame,
    QTableWidget, QTableWidgetItem, QComboBox, QTabWidget, QStackedWidget
)
import os
from train import GADF, generate_cwt, EMTF
from sklearn.preprocessing import MinMaxScaler
import scipy.io as sio
from PIL import Image
from PyQt5.QtCore import QThread, pyqtSignal
import numpy as np
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QStandardItemModel, QStandardItem
from PyQt5.QtWidgets import QAbstractItemView
from torchvision import transforms
from torch.utils.data import Dataset
from torch.utils.data import DataLoader, TensorDataset, Subset
import time
from PyQt5.QtGui import QPixmap  # PyQt5

import matplotlib.pyplot as plt
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtGui import QImage
from PyQt5.QtCore import QModelIndex
from pathlib import Path

'''############ 线程类处理数据 #############'''
class DataHandleThread(QThread):
    progress = pyqtSignal(str)  # 处理进度
    finished = pyqtSignal(dict, int)     # 处理完成信号

    def __init__(self, files, label_index):
        super().__init__()
        self.files = files
        self.label_index = label_index
        self.list = {}

        # 定义 transform
        self.transform = transforms.Compose([
            transforms.Resize((128, 128)),  # 改为标准尺寸
            transforms.RandomCrop(128),
            transforms.RandomVerticalFlip(),
            transforms.RandomRotation(20),  # 角度单位为度
            transforms.ColorJitter(brightness=0.1, contrast=0.1,
                                   saturation=0.1, hue=0.1),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485], std=[0.229])
        ])

    def run(self):
        self.progress.emit(f"[Label {self.label_index}] Start processing {len(self.files)} files...")
        for f in self.files:
            print(222)
            self.progress.emit(f"[Label {self.label_index}] Processing {f}...")
            self.data_handle(f)
        self.finished.emit(self.list, self.label_index)

    def data_handle(self, path_data):
        try:
            if f"GADF_{self.label_index}" not in self.list:
                self.list[f"GADF_{self.label_index}"] = []
            if f"CWT_{self.label_index}" not in self.list:
                self.list[f"CWT_{self.label_index}"] = []
            if f"EMTF_{self.label_index}" not in self.list:
                self.list[f"EMTF_{self.label_index}"] = []

            root_dir = Path(path_data)
            print("正在处理路径:", root_dir)

            cates = ["GADF", "CWT", "EMTF"]
            for cate in cates:
                cate_path = root_dir / cate
                print("检查目录:", cate_path, "存在吗？", cate_path.exists())
                for img_path in cate_path.glob("*.png"):
                    # print("找到图片:", img_path)
                    img = Image.open(img_path)
                    M = self.transform(img)
                    self.list[f"{cate}_{self.label_index}"].append(M)
                    print(f"{cate}_{self.label_index}")

            print("DataHandleThread线程完成")
        except Exception as e:
            print(f"Error processing: {e}")

class ModelInitThread(QThread):
    finished = pyqtSignal(dict, object)  # modell, loss_fn

    def __init__(self, model_names, lr, weight_decay, gamma):
        super().__init__()
        self.model_names = model_names
        self.lr = lr
        self.weight_decay = weight_decay
        self.gamma = gamma

    def run(self):
        try:
            modell, loss_fn = model_de_specific(self.model_names, self.lr, self.weight_decay, self.gamma)
            self.finished.emit(modell, loss_fn)
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("get_selected_models 线程报错：\n", error_msg)  # 控制台显示


class DataPrepareThread(QThread):
    finished = pyqtSignal(object, object, object)  # GADF_dataset, CWT_dataset, EMTF_dataset

    def __init__(self, data_dict):
        super().__init__()
        self.data_dict = data_dict

    def run(self):
        try:
            import torch
            # 合并所有标签的数据
            GADF_list, CWT_list, EMTF_list = [], [], []
            GADF_labels, CWT_labels, EMTF_labels = [], [], []
            for label_idx, sub_dict in self.data_dict.items():
                gadf_items = sub_dict.get(f"GADF_{label_idx}", [])
                cwt_items = sub_dict.get(f"CWT_{label_idx}", [])
                emtf_items = sub_dict.get(f"EMTF_{label_idx}", [])
                GADF_list.extend(gadf_items);CWT_list.extend(cwt_items);EMTF_list.extend(emtf_items)
                GADF_labels.extend([label_idx] * len(gadf_items))
                CWT_labels.extend([label_idx] * len(cwt_items))
                EMTF_labels.extend([label_idx] * len(emtf_items))

            print("124:GADF_label_unique:", np.unique(GADF_labels).tolist())
            print("125:CWT_label_unique:", np.unique(CWT_labels).tolist())
            print("126:EMTF_label_unique:", np.unique(EMTF_labels).tolist())

            # 创建 Dataset
            GADF_dataset = MyTensorDataset(GADF_list, GADF_labels)
            CWT_dataset = MyTensorDataset(CWT_list, CWT_labels)
            EMTF_dataset = MyTensorDataset(EMTF_list, EMTF_labels)

            # 打乱索引
            dataset_size = len(GADF_dataset)
            indices = torch.randperm(dataset_size)

            GADF_dataset = Subset(GADF_dataset, indices)
            CWT_dataset = Subset(CWT_dataset, indices)
            EMTF_dataset = Subset(EMTF_dataset, indices)

            BATCH_SIZE = 32
            GADF_loader = DataLoader(GADF_dataset, batch_size=BATCH_SIZE, shuffle=False)
            CWT_loader = DataLoader(CWT_dataset, batch_size=BATCH_SIZE, shuffle=False)
            EMTF_loader = DataLoader(EMTF_dataset, batch_size=BATCH_SIZE, shuffle=False)

            data_iter = iter(GADF_loader)
            batch_x, batch_y = next(data_iter)  # 取第一个 batch
            print("DataPrepareThread中一批次的标签", batch_y)

            self.finished.emit(GADF_loader, CWT_loader, EMTF_loader)
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("get_selected_models 线程报错：\n", error_msg)  # 控制台显示

class TrainThread(QThread):
    log_signal = pyqtSignal(str)          # 用于输出日志到GUI
    finished = pyqtSignal(dict)           # 训练完成，返回结果（准确率等）
    current_progress = pyqtSignal(str)    # 返回正在进行epoch
    learning_rate = pyqtSignal(str)       # 返回学习率
    extimated_remaining_time = pyqtSignal(str)  # 返回预估剩余时间
    acc_loss = pyqtSignal(bytes)

    def __init__(self, selected_models, modell, loss_fn, loaders, epochs=15):
        super().__init__()
        self.selected_models = selected_models
        self.modell = modell
        self.loss_fn = loss_fn
        self.loaders = loaders   # {'GADF': gadf_loader, 'CWT': cwt_loader, 'EMTF': emtf_loader}
        self.epochs = epochs
        self.acc_list = []
        self.loss_list = []
        self.al = {}

    def run(self):
        try:
            results = {}
            for name in self.selected_models:
                self.log_signal.emit(f"正在训练 {name} ...")
                print(f"正在训练 {name} ...")

                if name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                    acc = self.p2D_train(name, self.loss_fn, self.loaders['GADF'], self.loaders['EMTF'], self.modell,
                                         self.epochs)
                elif name == "P2D_DSResNet_CWT":
                    acc = self.train1(name, self.loss_fn, self.loaders['CWT'], self.modell, self.epochs)
                else:
                    acc = self.train1(name, self.loss_fn, self.loaders['GADF'], self.modell, self.epochs)

                results[name] = acc
                self.log_signal.emit(f"{name} 训练完成，准确率: {acc:.4f}")
            self.finished.emit(results)  # 返回所有模型的结果
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("get_selected_models 线程报错：\n", error_msg)  # 控制台显示

    def train1(self, name: str, loss_fn, train_dl, amodel: dict, epochs: int):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        start_time = time.perf_counter()
        name_model = amodel[name]
        model = name_model["model"]
        opt = name_model["opt"]
        scheduler = name_model["scheduler"]

        # 刚开始汇报一次学习率
        self.learning_rate.emit(f"{scheduler.get_last_lr()[0]:.6f}")

        epoch_acc = 0.0
        for epoch in range(epochs):
            self.current_progress.emit(f"{epoch + 1}/{epochs}")
            t0 = time.perf_counter()
            correct = 0
            total = 0
            running_loss = 0.0

            model.train()
            for x, y in train_dl:
                y = y.long()
                x, y = x.to(device), y.to(device)
                y_pred = model(x)
                loss = loss_fn(y_pred, y)
                opt.zero_grad()
                loss.backward()
                opt.step()

                with torch.no_grad():
                    pred = torch.argmax(y_pred, dim=1)
                    correct += (pred == y).sum().item()
                    total += y.size(0)
                    running_loss += loss.item() * y.size(0)

            scheduler.step()
            # 每个 epoch 汇报一次最新学习率
            self.learning_rate.emit(f"{scheduler.get_last_lr()[0]:.6f}")

            epoch_acc = correct / max(total, 1)
            epoch_loss = running_loss / max(total, 1)

            # 剩余时间估计
            elapsed = time.perf_counter() - t0
            rest_time = elapsed * (epochs - (epoch + 1))
            min_, sec_ = self.get_time(rest_time)
            self.extimated_remaining_time.emit(f"{int(min_)}:{int(sec_)}")

            # 追加当轮指标
            self.acc_list.append(round(epoch_acc, 4))
            self.loss_list.append(round(epoch_loss, 4))

            # ✅ 在每个 epoch 末尾生成并发送曲线 PNG
            fig_bytes = self.make_acc_loss_png_bytes(self.acc_list, self.loss_list, epochs_total=self.epochs)
            self.acc_loss.emit(fig_bytes)

            # 日志
            self.log_signal.emit(
                f"{name} Epoch [{epoch + 1}/{epochs}]  loss={epoch_loss:.4f}  acc={epoch_acc:.4f}  "
                f"ETA {int(min_)}m {int(sec_)}s"
            )

            # 按需保存权重
            save_path = r"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\model_save"
            torch.save(model.state_dict(), f"{save_path}\\{name}_best.pth")

        total_time = time.perf_counter() - start_time
        min_, sec_ = self.get_time(total_time)
        print(f"all_time: {int(min_)} minutes, {int(sec_)} seconds")
        return epoch_acc

    def p2D_train(self, name: str, loss_fn, GADF_train_dl, CWT_train_dl, amodel: dict, epochs: int):
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        start_time = time.perf_counter()

        name_model = amodel[name]
        model = name_model["model"].to(device)
        optim_ = name_model["opt"]
        scheduler = name_model["scheduler"]

        # 汇报初始学习率
        self.learning_rate.emit(f"{scheduler.get_last_lr()[0]:.6f}")

        epoch_acc = 0.0
        for epoch in range(epochs):
            self.current_progress.emit(f"{epoch + 1}/{epochs}")
            t0 = time.perf_counter()
            correct = 0
            total = 0
            running_loss = 0.0

            model.train()
            for (GADF_x, GADF_labels), (CWT_x, CWT_labels) in zip(GADF_train_dl, CWT_train_dl):
                y = GADF_labels.long()
                GADF_x, CWT_x, y = GADF_x.to(device), CWT_x.to(device), y.to(device)

                y_pred = model(GADF_x, CWT_x)
                loss = loss_fn(y_pred, y)

                optim_.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optim_.step()

                with torch.no_grad():
                    pred = torch.argmax(y_pred, dim=1)
                    correct += (pred == y).sum().item()
                    total += y.size(0)
                    running_loss += loss.item() * y.size(0)

            scheduler.step()
            self.learning_rate.emit(f"{scheduler.get_last_lr()[0]:.6f}")

            epoch_acc = correct / max(total, 1)
            epoch_loss = running_loss / max(total, 1)

            # 累加、画图、发信号（每个 epoch）
            self.acc_list.append(round(epoch_acc, 4))
            self.loss_list.append(round(epoch_loss, 4))
            fig_bytes = self.make_acc_loss_png_bytes(self.acc_list, self.loss_list, epochs_total=self.epochs)
            self.acc_loss.emit(fig_bytes)

            # ETA
            elapsed = time.perf_counter() - t0
            min_, sec_ = self.get_time(elapsed * (epochs - (epoch + 1)))
            self.extimated_remaining_time.emit(f"{int(min_)}:{int(sec_)}")

            # 日志
            self.log_signal.emit(
                f"{name} Epoch [{epoch + 1}/{epochs}]  loss={epoch_loss:.4f}  acc={epoch_acc:.4f}  "
                f"ETA {int(min_)}m {int(sec_)}s"
            )

            # 保存
            save_path = r"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\model_save"
            torch.save(model.state_dict(), f"{save_path}\\{name}_best.pth")

        total_time = time.perf_counter() - start_time
        min_, sec_ = self.get_time(total_time)
        print(f"总训练时间: {int(min_)} 分 {int(sec_)} 秒")
        return epoch_acc

    def get_time(self, time):
        if time > 60:
            min = int(time // 60)
            sec = int(time % 60)
            return min, sec
        else:
            return 0, time

    def make_acc_loss_png_bytes(self, acc_list, loss_list, epochs_total: int = None) -> bytes:
        """
        绘制训练过程中 Loss / Accuracy 曲线
        每次调用都会返回一张 PNG 图片（二进制流）
        第一次绘图时若提供 epochs_total，将固定 X 轴范围为 1~epochs_total
        """
        from matplotlib.figure import Figure
        from matplotlib.backends.backend_agg import FigureCanvasAgg
        from io import BytesIO
        import numpy as np

        if not acc_list and not loss_list:
            return b""

        # 当前已经训练的 epoch 数
        epochs_done = max(len(acc_list), len(loss_list))
        epochs = np.arange(1, epochs_done + 1)

        fig = Figure(figsize=(7, 4), dpi=120)
        ax1 = fig.add_subplot(111)
        ax2 = ax1.twinx()

        # ====== 绘制 Loss（红） ======
        if loss_list:
            loss_line, = ax1.plot(
                epochs[:len(loss_list)], loss_list,
                linewidth=2, marker='o', markersize=3,
                label='Loss', color='#D62728'  # 红
            )
        else:
            loss_line = None

        # ====== 绘制 Accuracy（蓝） ======
        if acc_list:
            acc_line, = ax2.plot(
                epochs[:len(acc_list)], acc_list,
                linewidth=2, marker='s', markersize=3,
                label='Accuracy', color='#1F77B4'  # 蓝
            )
        else:
            acc_line = None

        # ====== 高亮当前 epoch ======
        cur_epoch = epochs_done
        ax1.axvline(cur_epoch, linestyle='--', linewidth=1, color='#888888', alpha=0.8)
        if loss_list:
            ax1.scatter([cur_epoch], [loss_list[-1]], s=38, zorder=5, color='#D62728')
        if acc_list:
            ax2.scatter([cur_epoch], [acc_list[-1]], s=38, zorder=5, color='#1F77B4')

        # ====== 坐标轴设定 ======
        # 横轴范围：第一次调用时就固定为 1~epochs_total（若有传入）
        if epochs_total:
            ax1.set_xlim(1, epochs_total)
        else:
            ax1.set_xlim(1, epochs_done)

        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax2.set_ylabel('Accuracy')
        ax1.grid(alpha=0.28)

        # ====== 自动纵轴留白 ======
        if loss_list:
            ymin, ymax = float(np.min(loss_list)), float(np.max(loss_list))
            pad = (ymax - ymin) * 0.12 or 0.1
            ax1.set_ylim(ymin - pad, ymax + pad)
        if acc_list:
            ymin, ymax = float(np.min(acc_list)), float(np.max(acc_list))
            pad = (ymax - ymin) * 0.12 or 0.02
            ax2.set_ylim(max(0.0, ymin - pad), min(1.0, ymax + pad))

        # ====== 图例 ======
        lines = [l for l in (loss_line, acc_line) if l is not None]
        labels = [l.get_label() for l in lines]
        if lines:
            ax1.legend(lines, labels, loc='upper right', frameon=False)

        fig.tight_layout(pad=1.2)

        # ====== 输出 PNG bytes ======
        canvas = FigureCanvasAgg(fig)
        canvas.draw()
        buf = BytesIO()
        fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
        return buf.getvalue()


class TsneThread(QThread):
    log_signal = pyqtSignal(str)       # 输出日志
    finished = pyqtSignal(str)         # 完成后传回图片路径

    def __init__(self, name, model, dataloader, device, save_path):
        super().__init__()
        self.model = model
        self.name = name
        self.dataloader = dataloader
        self.device = device
        self.save_path = save_path

    def run(self):
        try:
            from tsne_function import (CNN_extract_features, Proto_extract_features, GNN_extract_features,
                                       DN4_extract_features,
                                       DSResNet_extract_features, DSResNet_GADF_extract_features,
                                       DSResNetCWT_extract_features,
                                       P2D_DSResNet_V2_extract_features, P2D_DSResNet_1_extract_features)
            from sklearn.manifold import TSNE

            self.log_signal.emit("开始提取特征...")
            if self.name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                if self.name == "P2D_GE_Net":
                    features, labels = DSResNet_extract_features(self.model, self.dataloader[0], self.dataloader[1], self.device)
                if self.name == "P2D_DSResNet_1":
                    features, labels = P2D_DSResNet_1_extract_features(self.model, self.dataloader[0], self.dataloader[1], self.device)
                if self.name == "P2D_DSResNet_V2":
                    features, labels = P2D_DSResNet_V2_extract_features(self.model, self.dataloader[0], self.dataloader[1], self.device)
            elif self.name == "P2D_DSResNet_CWT":
                features, labels = DSResNetCWT_extract_features(self.model, self.dataloader, self.device)
            else:
                if self.name == "CNNNet":
                    features, labels = CNN_extract_features(self.model, self.dataloader, self.device)
                if self.name == "DN4Net":
                    features, labels = DN4_extract_features(self.model, self.dataloader, self.device)
                if self.name == "GNNImageClassifier":
                    features, labels = GNN_extract_features(self.model, self.dataloader, self.device)
                if self.name == "DSResNet_GADF":
                    features, labels = DSResNet_GADF_extract_features(self.model, self.dataloader, self.device)
            self.log_signal.emit(f"提取完成，得到特征: {features.shape}")

            self.log_signal.emit("开始执行 T-SNE 降维...")
            tsne = TSNE(n_components=2, perplexity=20, n_iter=1000, random_state=42)
            features_2d = tsne.fit_transform(features)
            # --- t-SNE 绘图（高清） ---
            plt.figure(figsize=(10, 8), dpi=300)

            p_labels = ['nomal_data', 'inner_ring', 'Outer_ring', 'Rolling_element']
            classes = np.unique(labels)

            for cls in classes:
                idxs = labels == cls
                # 点更大、透明度略高、避免边线造成锯齿
                plt.scatter(
                    features_2d[idxs, 0],
                    features_2d[idxs, 1],
                    s=16, alpha=0.85, edgecolors='none', label=f"{p_labels[cls]}"
                )

            # 坐标轴与标题、图例的字号增大
            plt.xlabel("t-SNE dim 1", fontsize=14)
            plt.ylabel("t-SNE dim 2", fontsize=14)
            plt.title(f"t-SNE of {self.name} Features", fontsize=16, weight="bold")
            plt.grid(alpha=0.25)

            leg = plt.legend(title="Class Label", fontsize=12)
            if leg and leg.get_title():
                leg.get_title().set_fontsize(12)

            plt.tight_layout()
            # 关键：更高 dpi + 去掉多余留白
            plt.savefig(self.save_path, dpi=400, bbox_inches='tight')
            plt.close()
            self.log_signal.emit("T-SNE 绘制完成")
            self.finished.emit(self.save_path)

        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("TSNE 线程报错：\n", error_msg)  # 控制台显示


class ConfusionMatrixThread(QThread):
    log_signal = pyqtSignal(str)      # 日志输出
    finished = pyqtSignal(str)        # 完成后返回图片路径

    def __init__(self, model_name, model, dataloader, device, save_path):
        super().__init__()
        self.model_name = model_name
        self.model = model
        self.dataloader = dataloader
        self.device = device
        self.save_path = save_path

    def run(self):
        try:
            from sklearn.metrics import confusion_matrix
            import seaborn as sns
            import numpy as np
            import torch

            self.log_signal.emit(f"开始生成混淆矩阵: {self.model_name}")
            self.model.eval()
            all_preds = []
            all_labels = []

            with torch.no_grad():
                if self.model_name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                    for (G_x, G_y), (x, y) in zip(self.dataloader[0], self.dataloader[1]):
                        G_x, x, y = G_x.to(self.device), x.to(self.device), y.to(self.device)
                        outputs = self.model(G_x, x)
                        preds = torch.argmax(outputs, dim=1)
                        all_preds.append(preds.cpu())
                        all_labels.append(y.cpu())
                elif self.model_name == "P2D_DSResNet_CWT":
                    for x, y in self.dataloader:
                        x, y = x.to(self.device), y.to(self.device)
                        outputs = self.model(x)
                        preds = torch.argmax(outputs, dim=1)
                        all_preds.append(preds.cpu())
                        all_labels.append(y.cpu())
                else:
                    for x, y in self.dataloader:
                        x, y = x.to(self.device), y.to(self.device)
                        outputs = self.model(x)
                        preds = torch.argmax(outputs, dim=1)
                        all_preds.append(preds.cpu())
                        all_labels.append(y.cpu())

            all_preds = torch.cat(all_preds).numpy()
            all_labels = torch.cat(all_labels).numpy()

            cm = confusion_matrix(all_labels, all_preds)

            cm_normalized = cm.astype("float")/cm.sum(axis=1)[: , np.newaxis]

            plt.figure(figsize=(10, 8), dpi=300)

            sns.heatmap(
                cm_normalized,
                annot=True,
                fmt=".2f",
                cmap="Blues",
                cbar=True,
                annot_kws={"size": 12},  # 加大文字
                linewidths=0.3,  # 增强格线
                linecolor="white"  # 格线颜色
            )

            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)
            plt.xlabel("预测类别", fontsize=14)
            plt.ylabel("真实类别", fontsize=14)
            plt.title(f"{self.model_name} 混淆矩阵", fontsize=16, weight="bold")
            plt.tight_layout()
            plt.savefig(self.save_path, dpi=400, bbox_inches="tight")
            plt.close()

            self.log_signal.emit("混淆矩阵绘制完成")
            self.finished.emit(self.save_path)
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("ConfusionMatrixThread 线程报错：\n", error_msg)  # 控制台显示


# ===== 预测线程 =====
class PredictThread(QThread):
    log_signal = pyqtSignal(str)       # 日志信号
    result_signal = pyqtSignal(object)   # 预测结果 {class_name: prob, ...}
    # 把时域波形和频谱图以 PNG bytes 的形式发出去
    waveform_ready = pyqtSignal(bytes)
    spectrum_ready = pyqtSignal(bytes)
    final_advice = pyqtSignal(str, str)  # (label_name, advice_text)

    def __init__(self, model_name, model, filepath, device, fs=12000, seg_len=2048):
        super().__init__()
        self.model_name = model_name
        self.model = model
        self.filepath = filepath
        self.device = device
        self.fs = fs
        self.seg_len = seg_len
        # 预处理
        self.transform = transforms.Compose([
            transforms.Grayscale(num_output_channels=1),
            transforms.Resize((128, 128)),
            transforms.ToTensor()
        ])

    def run(self):
        try:
            import torch
            from torchvision import transforms
            from PIL import Image
            import numpy as np
            import matplotlib.pyplot as plt
            from io import BytesIO
            from numpy.fft import rfft, rfftfreq

            self.log_signal.emit(f"加载文件: {self.filepath}")

            fs = int(self.fs)  # 使用传入的采样率
            start_index = 5 * fs
            scaler = MinMaxScaler(feature_range=(-1, 1))
            mat_data_dict = sio.loadmat(self.filepath)

            fourth_key = list(mat_data_dict.keys())[3]
            DE_time = mat_data_dict[fourth_key]
            DE_time = scaler.fit_transform(DE_time.reshape(-1, 1)).flatten()[start_index:]

            # 使用传入的分段长度
            seg_len = min(int(self.seg_len), DE_time.shape[0])
            sig_seg = DE_time[:seg_len]

            # 时域波形
            t = np.arange(seg_len) / fs
            fig1 = plt.figure(figsize=(6, 2.6), dpi=150)
            ax1 = fig1.add_subplot(111)
            ax1.plot(t, sig_seg)
            ax1.set_xlabel("Time(s)")
            ax1.set_ylabel("Ampliture")
            ax1.set_title("Waveform")
            ax1.grid(alpha=0.3)
            fig1.tight_layout()
            buf1 = BytesIO()
            fig1.savefig(buf1, format='png', bbox_inches='tight')
            plt.close(fig1)
            self.waveform_ready.emit(buf1.getvalue())

            # 幅度图
            # 加 Hann 窗减少频谱泄漏
            win = np.hanning(seg_len)
            xw = sig_seg * win
            # 零均值（可选，避免直流分量过大）
            xw = xw - np.mean(xw)
            N = seg_len
            X = np.fft.fft(xw, n=N)
            freqs = np.fft.fftfreq(N, d=1.0 / fs)
            # 只保留非负频率部分
            pos_idx = freqs >= 0
            freqs = freqs[pos_idx]
            X = X[pos_idx]
            # 单边幅度谱归一化：|X|/N，再对非直流和非Nyquist分量乘2
            amp = np.abs(X) / N
            if N % 2 == 0:
                # 偶数点：索引0是DC，索引N/2是Nyquist
                if len(amp) > 2:
                    amp[1:-1] *= 2.0
            else:
                # 奇数点：索引0是DC，最高频不是单独Nyquist点
                if len(amp) > 1:
                    amp[1:] *= 2.0

            fig2 = plt.figure(figsize=(6, 2.6), dpi=150)
            ax2 = fig2.add_subplot(111)
            ax2.plot(freqs, amp)
            ax2.set_xlim(0, fs / 2)
            ax2.set_xlabel("Frequency (Hz)")
            ax2.set_ylabel("Amplitude")
            ax2.set_title("Amplitude Spectrum (FFT)")
            ax2.grid(alpha=0.3)
            fig2.tight_layout()
            buf2 = BytesIO()
            fig2.savefig(buf2, format='png', bbox_inches='tight')
            plt.close(fig2)
            self.spectrum_ready.emit(buf2.getvalue())

            phi_segments_line = np.clip(DE_time, -1.0, 1.0)
            phi_segments_line = np.arccos(phi_segments_line)
            if np.any(np.isnan(DE_time)) or np.any(np.isinf(DE_time)):
                print(f"Warning: DE_time contains NaN or Inf")
            '''###################修改'''
            window_size = min(1024, seg_len)  # 也可以用 seg_len
            step_size = window_size // 2
            num_windows = ((phi_segments_line.shape[0] - window_size) // step_size + 1)  # 向下取整
            count = 0
            GADF_IMG = []
            CWT_IMG = []
            EMTF_IMG = []
            pre_probs = []
            for num_window in range(num_windows):
                start = num_window * step_size
                end = start + window_size
                M_GADF = GADF(phi_segments_line[start: end], 1, count)
                M_CWT = generate_cwt(phi_segments_line[start: end], 1, count)
                M_GADF_tensor = self.transform(M_GADF).unsqueeze(0)
                M_CWT_tensor = self.transform(M_CWT).unsqueeze(0)
                # M_EMTF = EMTF(phi_segments_line[start: end], window_size, 1, count)
                M_EMTF = 1
                count = count + 1
                GADF_IMG.append(M_GADF_tensor)
                CWT_IMG.append(M_CWT_tensor)
                EMTF_IMG.append(M_EMTF)

            self.model.eval()
            with torch.no_grad():
                if self.model_name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                    for GADF_img, CWT_img in zip(GADF_IMG, CWT_IMG):
                        GADF_img, CWT_img = GADF_img.to(self.device), CWT_img.to(self.device)
                        outputs = self.model(GADF_img, CWT_img)
                        preds = torch.argmax(outputs, dim=1).cpu().numpy()  # 预测标签
                        pre_probs.extend(preds.tolist())
                elif self.model_name == "P2D_DSResNet_CWT":
                    for CWT_img in CWT_IMG:
                        CWT_img = CWT_img.to(self.device)
                        outputs = self.model(CWT_img)
                        preds = torch.argmax(outputs, dim=1).cpu().numpy()  # 预测标签
                        pre_probs.extend(preds.tolist())
                else:
                    for GADF_img in GADF_IMG:
                        GADF_img = GADF_img.to(self.device)
                        outputs = self.model(GADF_img)
                        preds = torch.argmax(outputs, dim=1).cpu().numpy()  # 预测标签
                        pre_probs.extend(preds.tolist())

            # 转换为字典输出
            pre_probs = [int(x) for x in pre_probs]
            print(pre_probs)  # 现在这句也更安全
            self.log_signal.emit("预测完成！")
            self.result_signal.emit(pre_probs)

            from collections import Counter

            # ====== 多数投票得到最终类别 ======
            if len(pre_probs) == 0:
                # 没有窗口预测，兜底
                label_id = -1
                label_name = "无法判定"
            else:
                # 计数最多者
                cnt = Counter(pre_probs)
                label_id = cnt.most_common(1)[0][0]

            # ====== 将label_id映射为中文类别名 ======
            # 注意：这里的索引顺序需与你的训练标签一致！
            # 你在 TSNE 可视化里用过：['nomal_data','inner_ring','Outer_ring','Rolling_element']
            # 相应中文：
            id2name = {
                0: "正常轴承",
                1: "内圈故障轴承",
                2: "外圈故障轴承",
                3: "滚动体故障轴承",
            }
            label_name = id2name.get(label_id, f"未知类别({label_id})")

            # ====== 生成整改建议 ======
            def get_advice(lid: int) -> str:
                if lid == 0:
                    return (
                        "状态判定：正常。\n"
                        "建议：\n"
                        "1) 保持当前工况与润滑策略，按点检周期继续监测振动/温度；\n"
                        "2) 校验传感器安装与数据链路，避免误报；\n"
                        "3) 建议保留本次频谱与波形作为基线，用于后续趋势对比。"
                    )
                if lid == 1:  # 内圈故障
                    return (
                        "状态判定：内圈故障疑似。\n"
                        "建议：\n"
                        "1) 检查轴与内圈过盈量、同轴度和联轴器对中，必要时重装配；\n"
                        "2) 复核载荷是否超标/冲击载荷是否频繁，必要时优化工况；\n"
                        "3) 复核润滑：更换同品牌同型号润滑脂/油，确保粘度等级匹配转速；\n"
                        "4) 观察 BPFI 条纹/边频：若随转速/载荷增强，建议计划停机更换轴承并检查轴颈磨损；\n"
                        "5) 排查电蚀风险（变频器接地/轴电流），必要时加绝缘轴承或接地环。"
                    )
                if lid == 2:  # 外圈故障
                    return (
                        "状态判定：外圈故障疑似。\n"
                        "建议：\n"
                        "1) 检查座孔配合与椭圆度、松动与不良定位（轴承外圈是否转动/打滑）；\n"
                        "2) 复核基础/机座刚度与安装应力，确认没有装偏或敲击伤；\n"
                        "3) 排查污染与压痕（硬颗粒/轴向通道损伤），必要时更换密封并加装过滤；\n"
                        "4) 关注 BPFO 分量及其边频：若持续升高，安排计划检修并更换轴承；\n"
                        "5) 检查轴承座温升与振动趋势，避免带病长时间运行。"
                    )
                if lid == 3:  # 滚动体故障
                    return (
                        "状态判定：滚动体故障疑似。\n"
                        "建议：\n"
                        "1) 优先排查润滑不足/润滑脂劣化：补脂或换脂，并检查补脂周期与用量；\n"
                        "2) 检查密封与环境清洁度，防止异物进入导致点蚀与剥落；\n"
                        "3) 复核材料/热处理与超速运行风险，必要时更换更高等级轴承；\n"
                        "4) 关注球通过频率（BSF）及调制定子频（Fmod）：若伴随宽带噪声上升，尽快计划性更换；\n"
                        "5) 运行中降低冲击载荷与过度振动源（齿轮啮合不良、共振）。"
                    )
                # 未知/无法判定
                return (
                    "状态判定：信号不足或分类不确定。\n"
                    "建议：\n"
                    "1) 复查传感器安装与量程、采样频率与前置放大设置；\n"
                    "2) 采集更长时长与多工况（不同转速/载荷）的数据；\n"
                    "3) 同步温度、转速与油液监测做多模态交叉验证。"
                )

            advice = get_advice(label_id)

            # 发信号给主线程
            self.final_advice.emit(label_name, advice)

        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("PredictThread 线程报错：\n", error_msg)  # 控制台显示


def model_de_specific(model_names, lr=1e-3, weight_decay=1e-5, gamma=0.9):
    try:
        import torch
        import torch.nn as nn
        from Neural_Networks import (P2D_GE_Net, CNNNet, DN4Net, ProtoNetClassifier, GNNImageClassifier,
                                     DSResNet_GADF, P2D_DSResNet_CWT, P2D_DSResNet_1, P2D_DSResNet_V2)

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        loss_fn = nn.CrossEntropyLoss()
        modell = {}
        print("model_de_specific")

        def pack(model):
            opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
            scheduler = torch.optim.lr_scheduler.ExponentialLR(opt, gamma=gamma)
            return {"model": model, "opt": opt, "scheduler": scheduler}

        if "CNNNet" in model_names:
            modell["CNNNet"] = pack(CNNNet().to(device))
        if "ProtoNetClassifier" in model_names:
            modell["ProtoNetClassifier"] = pack(ProtoNetClassifier(in_channels=1, hidden_size=64, num_classes=4).to(device))
        if "DN4Net" in model_names:
            modell["DN4Net"] = pack(DN4Net(in_channels=1, num_classes=4).to(device))
        if "GNNImageClassifier" in model_names:
            modell["GNNImageClassifier"] = pack(GNNImageClassifier(in_channels=1, num_classes=4).to(device))
        if "P2D_GE_Net" in model_names:
            modell["P2D_GE_Net"] = pack(P2D_GE_Net(num_classes=4).to(device))
        if "DSResNet_GADF" in model_names:
            modell["DSResNet_GADF"] = pack(DSResNet_GADF(num_classes=4).to(device))
        if "P2D_DSResNet_CWT" in model_names:
            modell["P2D_DSResNet_CWT"] = pack(P2D_DSResNet_CWT(num_classes=4).to(device))
        if "P2D_DSResNet_1" in model_names:
            modell["P2D_DSResNet_1"] = pack(P2D_DSResNet_1(num_classes=4).to(device))
        if "P2D_DSResNet_V2" in model_names:
            modell["P2D_DSResNet_V2"] = pack(P2D_DSResNet_V2(num_classes=4).to(device))

        return modell, loss_fn
    except Exception as e:
        import traceback
        print("model_de_specific 线程报错：\n", traceback.format_exc())

def tensor_to_qpixmap(tensor):
    """tensor 转 QPixmap"""
    if isinstance(tensor, torch.Tensor):
        arr = tensor.detach().cpu().numpy()
    else:
        arr = tensor

    # 如果是 [C,H,W] 格式，转为 [H,W,C]
    if arr.ndim == 3 and arr.shape[0] in (1, 3):
        arr = arr.transpose(1, 2, 0)

    # 转成 0-255 uint8
    if arr.max() <= 1.0:
        arr = (arr * 255).astype(np.uint8)
    else:
        arr = arr.astype(np.uint8)

    h, w = arr.shape[:2]
    if arr.ndim == 2:  # 灰度
        q_image = QImage(arr.data, w, h, w * 3, QImage.Format_RGB888)
    else:  # RGB
        q_image = QImage(arr.data, w, h, w, QImage.Format_Grayscale8)

    return QPixmap.fromImage(q_image)


class ImageBrowser(QWidget):
    """一个小部件：支持左右翻页浏览一组图片"""
    def __init__(self, images):
        super().__init__()
        self.images = images

        layout = QVBoxLayout()

        # 图片显示区域
        self.stacked = QStackedWidget()
        for img in self.images:
            lbl = QLabel()
            lbl.setScaledContents(True)
            lbl.setPixmap(tensor_to_qpixmap(img))
            self.stacked.addWidget(lbl)
        layout.addWidget(self.stacked)

        # 控制按钮
        btn_layout = QHBoxLayout()
        self.btn_prev = QPushButton("上一张")
        self.btn_next = QPushButton("下一张")
        btn_layout.addWidget(self.btn_prev)
        btn_layout.addWidget(self.btn_next)
        layout.addLayout(btn_layout)

        self.setLayout(layout)

        # 绑定按钮
        self.btn_prev.clicked.connect(self.show_prev)
        self.btn_next.clicked.connect(self.show_next)

    def show_prev(self):
        idx = (self.stacked.currentIndex() - 1) % self.stacked.count()
        self.stacked.setCurrentIndex(idx)

    def show_next(self):
        idx = (self.stacked.currentIndex() + 1) % self.stacked.count()
        self.stacked.setCurrentIndex(idx)


class SecondWindow(QWidget):
    def __init__(self, data_dict):
        super().__init__()
        try:
            self.setWindowTitle("新界面")
            self.resize(1200, 1000)
            self.data_dict = data_dict

            layout = QVBoxLayout()
            self.tabs = QTabWidget()
            layout.addWidget(self.tabs)
            self.setLayout(layout)

            sample_key = list(self.data_dict.keys())[0]
            methods = self.data_dict[sample_key].keys()

            for method in methods:
                method = method[:-2]

                tab = QWidget()
                vbox = QVBoxLayout()

                images = []
                for k in self.data_dict.keys():
                    print(self.data_dict[k][f"{method}_{k}"])
                    images.extend(self.data_dict[k][f"{method}_{k}"])

                browser = ImageBrowser(images)
                vbox.addWidget(browser)
                tab.setLayout(vbox)
                self.tabs.addTab(tab, method)
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("SecondWindow 线程报错：\n", error_msg)  # 控制台显示


class EvalPanel(QWidget):
    def __init__(self, modell=None, loaders=None, tsne_max_model_combo=None, parent=None):
        super().__init__(parent)
        self.modell = modell or {}
        self.loaders = loaders or {}
        self.tsne_max_model_combo = tsne_max_model_combo or []

        from PyQt5.QtWidgets import QSizePolicy

        main_layout = QVBoxLayout(self)

        # ===== 顶层：三列（左：控制区，中：混淆矩阵图，右：T-SNE图） =====
        content_layout = QHBoxLayout()

        # ========== 左列：控制区（竖直） ==========
        controls_layout = QVBoxLayout()

        # ---- 混淆矩阵控制块 ----
        self.confusion_label = QLabel("混淆矩阵")
        self.confusion_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        controls_layout.addWidget(self.confusion_label)

        self.confusion_model_combo = QComboBox()
        self.confusion_model_combo.addItem("请选择模型")
        if self.tsne_max_model_combo:
            self.confusion_model_combo.addItems(self.tsne_max_model_combo)
        controls_layout.addWidget(self.confusion_model_combo)

        self.confusion_start_btn = QPushButton("开始")
        self.confusion_start_btn.clicked.connect(self.run_confusion_matrix)
        controls_layout.addWidget(self.confusion_start_btn)

        self.confusion_save_btn = QPushButton("保存图片")
        self.confusion_save_btn.clicked.connect(self.save_confusion_img)
        controls_layout.addWidget(self.confusion_save_btn)

        # 分隔线
        sep = QFrame()
        sep.setFrameShape(QFrame.HLine)
        sep.setFrameShadow(QFrame.Sunken)
        controls_layout.addWidget(sep)

        # ---- T-SNE 控制块 ----
        self.tsne_label = QLabel("T-SNE")
        self.tsne_label.setStyleSheet("font-weight: bold; font-size: 14px;")
        controls_layout.addWidget(self.tsne_label)

        self.tsne_model_combo = QComboBox()
        self.tsne_model_combo.addItem("请选择模型")
        if self.tsne_max_model_combo:
            self.tsne_model_combo.addItems(self.tsne_max_model_combo)
        controls_layout.addWidget(self.tsne_model_combo)

        self.tsne_start_btn = QPushButton("开始")
        self.tsne_start_btn.clicked.connect(self.run_tsne)
        controls_layout.addWidget(self.tsne_start_btn)

        self.tsne_save_btn = QPushButton("保存图片")
        self.tsne_save_btn.clicked.connect(self.save_tsne_img)
        controls_layout.addWidget(self.tsne_save_btn)

        controls_layout.addStretch(1)
        content_layout.addLayout(controls_layout, stretch=0)

        # ========== 中列：混淆矩阵图片区域 ==========
        self.confusion_area = QLabel()
        self.confusion_area.setMinimumSize(400, 300)
        self.confusion_area.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.confusion_area.setFrameStyle(QFrame.Box | QFrame.Raised)
        self.confusion_area.setAlignment(Qt.AlignCenter)
        self.confusion_area.setText("混淆矩阵图片区域")
        self.confusion_area.setFixedSize(500, 375)
        content_layout.addWidget(self.confusion_area, alignment=Qt.AlignCenter)

        # ========== 右列：T-SNE 图片区域 ==========
        self.tsne_area = QLabel()
        self.tsne_area.setMinimumSize(400, 300)
        self.tsne_area.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.tsne_area.setFrameStyle(QFrame.Box | QFrame.Raised)
        self.tsne_area.setAlignment(Qt.AlignCenter)
        self.tsne_area.setText("T-SNE 图片区域")
        self.tsne_area.setFixedSize(500, 375)
        content_layout.addWidget(self.tsne_area, alignment=Qt.AlignCenter)

        # 装配主布局
        main_layout.addLayout(content_layout, stretch=1)

    def update_context(self, modell, loaders, tsne_max_model_combo):
        """更新上下文（模型、数据加载器、下拉框）"""
        self.modell = modell
        self.loaders = loaders
        self.tsne_max_model_combo = tsne_max_model_combo or []

        # 刷新下拉框内容
        self.confusion_model_combo.clear()
        self.tsne_model_combo.clear()
        self.confusion_model_combo.addItem("请选择模型")
        self.tsne_model_combo.addItem("请选择模型")
        if self.tsne_max_model_combo:
            self.confusion_model_combo.addItems(self.tsne_max_model_combo)
            self.tsne_model_combo.addItems(self.tsne_max_model_combo)

    # ======= T-SNE 可视化 =======
    def run_tsne(self):
        try:
            model_name = self.tsne_model_combo.currentText()
            if model_name == "请选择模型":
                QMessageBox.warning(self, "提示", "请先选择一个模型！")
                return
            model = self.modell[model_name]["model"]

            if model_name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                train_dl = [self.loaders['GADF'], self.loaders['CWT']]
            elif model_name == "P2D_DSResNet_CWT":
                train_dl = self.loaders['CWT']
            else:
                train_dl = self.loaders['GADF']

            model.eval()
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            save_dir = "./tsne_results"
            os.makedirs(save_dir, exist_ok=True)
            tsne_img_path = os.path.join(save_dir, f"{model_name}_temp_tsne.png")

            self.tsne_thread = TsneThread(model_name, model, train_dl, self.device, tsne_img_path)
            self.tsne_thread.finished.connect(self._show_tsne_result)
            self.tsne_thread.start()

        except Exception:
            import traceback
            print("run_tsne 报错：\n", traceback.format_exc())

    def _show_tsne_result(self, img_path):
        """显示 T-SNE 结果图片"""
        pixmap = QPixmap(img_path).scaled(
            self.tsne_area.width(), self.tsne_area.height(),
            Qt.KeepAspectRatio, Qt.SmoothTransformation
        )
        self.tsne_area.setPixmap(pixmap)
        QMessageBox.information(self, "完成", "T-SNE 可视化完成！")

    def save_tsne_img(self):
        """保存 T-SNE 图片"""
        if self.tsne_area.pixmap() is None:
            QMessageBox.warning(self, "提示", "当前没有可保存的图片！")
            return
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存图片", "tsne.png", "PNG Files (*.png);;JPG Files (*.jpg)"
        )
        if file_path:
            self.tsne_area.pixmap().save(file_path)
            QMessageBox.information(self, "成功", f"图片已保存到 {file_path}")

    # ======= 混淆矩阵 =======
    def run_confusion_matrix(self):
        try:
            model_name = self.confusion_model_combo.currentText()
            if model_name == "请选择模型":
                QMessageBox.warning(self, "提示", "请先选择模型！")
                return

            model = self.modell[model_name]["model"]

            if model_name in ["P2D_GE_Net", "P2D_DSResNet_1", "P2D_DSResNet_V2"]:
                train_dl = [self.loaders['GADF'], self.loaders['CWT']]
            elif model_name == "P2D_DSResNet_CWT":
                train_dl = self.loaders['CWT']
            else:
                train_dl = self.loaders['GADF']

            save_dir = "./ConfusionMatrix_results"
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, f"{model_name}_ConfusionMatrix.png")
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            self.cm_thread = ConfusionMatrixThread(model_name, model, train_dl, device, save_path)
            self.cm_thread.finished.connect(self._show_confusion_result)
            self.cm_thread.start()

        except Exception:
            import traceback
            print("run_confusion_matrix 报错：\n", traceback.format_exc())

    def _show_confusion_result(self, img_path):
        """显示混淆矩阵结果图片"""
        pixmap = QPixmap(img_path)
        self.confusion_area.setPixmap(
            pixmap.scaled(self.confusion_area.size(), Qt.KeepAspectRatio)
        )
        QMessageBox.information(self, "完成", "混淆矩阵生成完成！")

    def save_confusion_img(self):
        """保存混淆矩阵图片"""
        if not self.confusion_area.pixmap():
            QMessageBox.warning(self, "提示", "没有图片可以保存！")
            return
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存混淆矩阵图片", "",
            "PNG Files (*.png);;All Files (*)"
        )
        if file_path:
            self.confusion_area.pixmap().save(file_path)

class PreprocessArea(QWidget):
    preprocessedReady = pyqtSignal(dict)  # 传出预处理数据

    def __init__(self, label_items):
        super().__init__()
        # 不再创建任何UI，只保存传入回调/容器
        self.label_items = label_items
        self.threads = []
        self.results = {}
        # 可选：占位最小尺寸，避免布局时为0
        self.setFixedHeight(1)

    # open_new_window 不再需要（可删除）
    # def open_new_window(self):
    #     pass

    def start_preprocess(self):
        try:
            self.pending_threads = 0
            label_items = self.label_items()
            self.pending_threads = len(label_items)
            for i, item in enumerate(label_items):
                thread = DataHandleThread(item.files, i)
                self.threads.append(thread)
                thread.progress.connect(lambda msg: print(msg))

                def handle_finished(result, label_idx=i):
                    try:
                        self.results[label_idx] = result
                        self.pending_threads -= 1
                        if self.pending_threads == 0:
                            # 全部完成后发信号
                            self.preprocessedReady.emit(self.results)
                    except Exception:
                        import traceback
                        print("handle_finished 报错：\n", traceback.format_exc())

                thread.finished.connect(handle_finished)
                thread.start()
        except Exception:
            import traceback
            print("start_preprocess 报错：\n", traceback.format_exc())


class CheckableComboBox(QComboBox):
    checkedChanged = pyqtSignal(list)  # ✅ 定义信号，放在这里
    def __init__(self, parent=None):
        super().__init__(parent)
        try:
            self.setModel(QStandardItemModel(self))
            # 不允许选择高亮项，防止选中后自动收起
            self.view().setSelectionMode(QAbstractItemView.NoSelection)
            # 用 pressed 信号手动切换勾选状态（不会触发收起）
            self.view().pressed.connect(self._toggle_check)
            # 显示已选项
            self.setEditable(True)
            self.lineEdit().setReadOnly(True)
            self.lineEdit().setPlaceholderText("请选择模型…")
        except Exception as e:
            print("出现错误:", e)

    def add_check_items(self, texts):
        for t in texts:
            it = QStandardItem(t)
            it.setFlags(Qt.ItemIsEnabled | Qt.ItemIsUserCheckable)
            it.setData(Qt.Unchecked, Qt.CheckStateRole)
            self.model().appendRow(it)

    def _toggle_check(self, index: QModelIndex):
        try:
            it = self.model().itemFromIndex(index)
            if it is not None:
                it.setCheckState(Qt.Unchecked if it.checkState() == Qt.Checked else Qt.Checked)
                self._refresh_text()
                self.checkedChanged.emit(self.checked_texts())
        except Exception as e:
            import traceback
            print("CheckableComboBox _toggle_check 报错:\n", traceback.format_exc())

    def _refresh_text(self):
        items = []
        for i in range(self.model().rowCount()):
            it = self.model().item(i)
            if it.checkState() == Qt.Checked:
                items.append(it.text())
        self.lineEdit().setText(", ".join(items))

    def checked_texts(self):
        out = []
        for i in range(self.model().rowCount()):
            it = self.model().item(i)
            if it.checkState() == Qt.Checked:
                out.append(it.text())
        return out


class LabelItem(QWidget):
    def __init__(self, index):
        super().__init__()
        self.index = index
        self.files = []

        layout = QVBoxLayout(self)

        # 上排：编号 + name输入框 + 加载数据按钮
        top_layout = QHBoxLayout()
        top_layout.setSpacing(5)
        top_layout.setContentsMargins(0, 0, 0, 0)

        self.label_id = QLabel(f"标签{index}:")

        self.label_name = QLineEdit()
        self.label_name.setPlaceholderText("name")
        self.label_name.setFixedWidth(120)  # 设置固定宽度，约能容纳11个字符

        self.btn_load = QPushButton("加载数据")
        self.btn_load.setFixedWidth(80)
        self.btn_load.clicked.connect(self.load_files)

        self.btn_reload = QPushButton("重新加载")
        self.btn_reload.setFixedWidth(80)
        self.btn_reload.clicked.connect(self.reload_files)

        top_layout.addWidget(self.label_id)
        top_layout.addWidget(self.label_name)
        top_layout.addWidget(self.btn_load)
        top_layout.addWidget(self.btn_reload)
        top_layout.addStretch(1)

        # 下排：显示文件名
        self.file_box = QTextEdit()
        self.file_box.setReadOnly(True)
        self.file_box.setFixedHeight(40)
        self.file_box.setFixedWidth(400)
        self.file_box.setStyleSheet("QTextEdit { background: #f0f0f0; }")

        layout.addLayout(top_layout)
        layout.addWidget(self.file_box)

    def load_files(self):
        folder = QFileDialog.getExistingDirectory(self, "选择文件夹")
        if folder:
            self.files = [folder]  # 注意，这里保存成列表
            folder_name = os.path.basename(folder)
            self.file_box.setText(folder_name)
            print("保存路径：", self.files)

    def reload_files(self):
        """清空文件名显示框和文件列表"""
        self.files = []
        self.file_box.clear()

class MyTensorDataset(Dataset):
    def __init__(self, samples, labels):
        self.samples = samples
        self.labels = labels

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        return self.samples[idx], self.labels[idx]

class TrainArea(QWidget):
    def __init__(self):
        super().__init__()
        try:
            self.preprocessed_data = {}  # 存储从 PreprocessArea 传过来的数据
            self.tsne_max_model_combo = []
            main_layout = QVBoxLayout(self)
            main_layout.setContentsMargins(10, 10, 10, 10)
            main_layout.setSpacing(10)
            # --------------------------------- 上方：数据集准备 + 训练区域（并排） ---------------------------------
            train_section = QHBoxLayout()

            # ===== 左侧：数据集准备 =====
            left_widget = QWidget()
            left_layout = QVBoxLayout(left_widget)
            left_layout.setContentsMargins(0, 0, 0, 0)
            left_layout.setSpacing(5)

            self.dataset_label = QLabel("数据集准备")
            self.dataset_label.setFixedHeight(30)
            left_layout.addWidget(self.dataset_label)

            label_layout = QHBoxLayout()
            label_layout.setSpacing(5)
            self.label_count_label = QLabel("标签数量:")

            self.label_count_input = QSpinBox()
            self.label_count_input.setRange(1, 20)
            self.label_count_input.setValue(4)
            self.label_count_input.valueChanged.connect(self.update_labels)

            label_layout.addWidget(self.label_count_label)
            label_layout.addWidget(self.label_count_input)

            self.btn_preprocess = QPushButton("开始预处理")
            self.btn_preprocess.clicked.connect(self.start_preprocess_clicked)
            label_layout.addWidget(self.btn_preprocess)

            label_layout.addStretch(1)
            left_layout.addLayout(label_layout)

            self.scroll_area = QScrollArea()
            self.scroll_area.setWidgetResizable(True)
            self.scroll_widget = QWidget()
            self.scroll_layout = QVBoxLayout(self.scroll_widget)
            self.scroll_layout.setContentsMargins(0, 0, 0, 0)
            self.scroll_layout.setSpacing(5)
            self.scroll_area.setWidget(self.scroll_widget)
            left_layout.addWidget(self.scroll_area)

            self.label_items = []
            self.update_labels()

            # ===== 保留无UI的预处理对象，并连接信号 =====
            self.preprocessed_area = PreprocessArea(lambda: self.label_items)
            self.preprocessed_area.preprocessedReady.connect(self.set_preprocessed_data)

            left_widget.setFixedWidth(480)
            train_section.addWidget(left_widget, 0)  # 左侧较窄

            # ===== 右侧：训练面板（含训练参数面板） =====
            right_col = QHBoxLayout()

            # —— 顶部：训练控制（选择模型 / 参数 / Train 按钮 / 输出）
            training_widget = QWidget()
            training_layout = QVBoxLayout(training_widget)
            training_layout.setSpacing(10)

            self.model_select = CheckableComboBox()
            models = [
                "CNNNet",
                "ProtoNetClassifier",
                "DN4Net",
                "GNNImageClassifier",
                "P2D_GE_Net",
                "DSResNet_GADF",
                "P2D_DSResNet_CWT",
                "P2D_DSResNet_1",
                "P2D_DSResNet_V2",
            ]
            self.model_select.add_check_items(models)
            self.model_select.checkedChanged.connect(
                lambda items: self.output_box.append("选择的模型: " + (", ".join(items) if items else "（未选择）"))
            )

            from PyQt5.QtWidgets import QGroupBox, QGridLayout, QDoubleSpinBox, QFormLayout

            param_group = QGroupBox("训练参数")
            param_grid = QGridLayout(param_group)

            # -------- Epochs --------
            self.epoch_spin = QSpinBox()
            self.epoch_spin.setRange(1, 2000)
            self.epoch_spin.setValue(15)

            # -------- Learning Rate --------
            self.lr_spin = QDoubleSpinBox()
            self.lr_spin.setDecimals(6)
            self.lr_spin.setRange(1e-7, 1.0)
            self.lr_spin.setSingleStep(1e-4)
            self.lr_spin.setValue(1e-3)

            # -------- Weight Decay --------
            self.wd_spin = QDoubleSpinBox()
            self.wd_spin.setDecimals(8)
            self.wd_spin.setRange(0.0, 1.0)
            self.wd_spin.setSingleStep(1e-6)
            self.wd_spin.setValue(1e-5)

            # -------- LR Gamma --------
            self.gamma_spin = QDoubleSpinBox()
            self.gamma_spin.setDecimals(3)
            self.gamma_spin.setRange(0.1, 0.999)
            self.gamma_spin.setSingleStep(0.01)
            self.gamma_spin.setValue(0.9)

            # ====== 布局（两行两列）======
            # 第 1 行：Epochs | Learning Rate
            param_grid.addWidget(QLabel("Epochs："), 0, 0)
            param_grid.addWidget(self.epoch_spin, 0, 1)
            param_grid.addWidget(QLabel("Learning Rate："), 0, 2)
            param_grid.addWidget(self.lr_spin, 0, 3)

            # 第 2 行：Weight Decay | LR Gamma
            param_grid.addWidget(QLabel("Weight Decay："), 1, 0)
            param_grid.addWidget(self.wd_spin, 1, 1)
            param_grid.addWidget(QLabel("LR Gamma："), 1, 2)
            param_grid.addWidget(self.gamma_spin, 1, 3)

            param_grid.setColumnStretch(1, 1)
            param_grid.setColumnStretch(3, 1)

            self.btn_train = QPushButton("Train")
            self.btn_train.clicked.connect(self.start_training)

            self.output_box = QTextEdit()
            self.output_box.setReadOnly(True)
            self.output_box.setFixedHeight(100)

            training_layout.addWidget(QLabel("选择模型:"))
            training_layout.addWidget(self.model_select)
            training_layout.addWidget(param_group)  # ✅ 训练参数面板在训练区域内
            training_layout.addWidget(self.btn_train)
            training_layout.addWidget(QLabel("输出:"))
            training_layout.addWidget(self.output_box)

            right_col.addWidget(training_widget, 0)

            # —— 中部：训练状态（进度/学习率/ETA）+ 曲线
            all_status_layout = QVBoxLayout()

            status_layout = QHBoxLayout()
            self.label_current_progress = QLabel("当前进度：")
            self.text_current_progress = QTextEdit()
            self.text_current_progress.setReadOnly(True)
            self.text_current_progress.setFixedSize(100, 40)
            self.center_text_edit(self.text_current_progress, "0")

            self.label_current_learning_rate = QLabel("当前学习率：")
            self.text_current_learning_rate = QTextEdit()
            self.text_current_learning_rate.setReadOnly(True)
            self.text_current_learning_rate.setFixedSize(100, 40)
            self.center_text_edit(self.text_current_learning_rate, "0.000")

            self.label_estimated_remaining_time = QLabel("预估剩余时间")
            self.text_estimated_remaining_time = QTextEdit()
            self.text_estimated_remaining_time.setReadOnly(True)
            self.text_estimated_remaining_time.setFixedSize(180, 40)
            self.center_text_edit(self.text_estimated_remaining_time, "00:00")

            status_layout.addWidget(self.label_current_progress)
            status_layout.addWidget(self.text_current_progress)
            status_layout.addWidget(self.label_current_learning_rate)
            status_layout.addWidget(self.text_current_learning_rate)
            status_layout.addWidget(self.label_estimated_remaining_time)
            status_layout.addWidget(self.text_estimated_remaining_time)
            all_status_layout.addLayout(status_layout)

            self.label_SQ_Line = QLabel("损失值/准确率曲线")
            self.SQ_area = QLabel()
            self.SQ_area.setFixedSize(500, 300)
            self.SQ_area.setFrameStyle(QFrame.Box | QFrame.Raised)
            self.SQ_area.setAlignment(Qt.AlignCenter)
            self.SQ_area.setText("损失值/准确率曲线")
            all_status_layout.addWidget(self.SQ_area, alignment=Qt.AlignCenter)

            right_col.addLayout(all_status_layout)

            # 右侧整列加入并排区域
            right_wrap = QWidget()
            right_wrap.setLayout(right_col)
            train_section.addWidget(right_wrap)  # 右侧更宽

            train_widget = QWidget()
            train_widget.setLayout(train_section)
            train_widget.setFixedHeight(420)  # 固定为 320 像素高
            # 先把并排的“训练区域”加入主布局（放在前面）
            main_layout.addWidget(train_widget, stretch=1)

            # --------------------------------- 评估面板（置于训练区域之后） ---------------------------------
            self.eval_panel = EvalPanel(self.modell if hasattr(self, 'modell') else {},
                                        getattr(self, 'loaders', {}),
                                        self.tsne_max_model_combo)
            main_layout.addWidget(self.eval_panel, stretch=4)

            # ========================= 预测区域 =========================
            bottom_widget = QWidget()
            bottom_layout = QVBoxLayout(bottom_widget)

            # ========== 上层：左(预测设置) 右(可视化) ==========
            upper_layout = QHBoxLayout()

            # ---- 上左：预测设置（参数、选择文件、模型、开始预测）
            settings_group = QGroupBox("预测设置")
            settings_group.setMinimumWidth(280)
            settings_v = QVBoxLayout(settings_group)

            # 文件选择
            self.select_file_btn = QPushButton("选择文件")
            self.select_file_btn.clicked.connect(self.select_file)
            self.select_file_btn.setFixedWidth(220)
            settings_v.addWidget(self.select_file_btn)

            # 模型选择
            self.model_combo = QComboBox()
            self.model_combo.addItem("请选择模型")
            self.model_combo.setFixedWidth(220)
            settings_v.addWidget(self.model_combo)

            # 预测参数
            param_group = QGroupBox("预测参数")
            param_form = QFormLayout(param_group)

            self.sr_spin = QDoubleSpinBox()
            self.sr_spin.setRange(1000, 200000)
            self.sr_spin.setDecimals(0)
            self.sr_spin.setSingleStep(1000)
            self.sr_spin.setValue(12000)
            self.sr_spin.setSuffix(" Hz")
            self.sr_spin.setFixedWidth(220)
            param_form.addRow("采样率：", self.sr_spin)

            self.seg_spin = QSpinBox()
            self.seg_spin.setRange(256, 32768)
            self.seg_spin.setSingleStep(256)
            self.seg_spin.setValue(1024)
            self.seg_spin.setFixedWidth(220)
            param_form.addRow("分段长度：", self.seg_spin)

            settings_v.addWidget(param_group)

            # 开始预测
            self.start_predict_btn = QPushButton("确定，开始预测")
            self.start_predict_btn.clicked.connect(self.start_predict)
            self.start_predict_btn.setFixedWidth(220)
            settings_v.addWidget(self.start_predict_btn)

            settings_v.addStretch(1)
            upper_layout.addWidget(settings_group, 0, Qt.AlignTop)

            # ---- 上右：原信号可视化（时域波形 + 频谱图）
            viz_right = QWidget()
            viz_right_h = QHBoxLayout(viz_right)

            # 时域波形
            waveform_layout = QVBoxLayout()
            self.waveform_label = QLabel("原信号时域波形")
            waveform_layout.addWidget(self.waveform_label)
            self.waveform_canvas = QLabel()
            self.waveform_canvas.setFixedSize(520, 260)
            self.waveform_canvas.setFrameStyle(QFrame.Box | QFrame.Raised)
            self.waveform_canvas.setAlignment(Qt.AlignCenter)
            self.waveform_canvas.setText("时域波形图区域")
            waveform_layout.addWidget(self.waveform_canvas, alignment=Qt.AlignCenter)
            viz_right_h.addLayout(waveform_layout)

            # 频谱图
            spectrum_layout = QVBoxLayout()
            self.spectrum_label = QLabel("原信号频域幅度谱")
            spectrum_layout.addWidget(self.spectrum_label)
            self.spectrum_canvas = QLabel()
            self.spectrum_canvas.setFixedSize(520, 260)
            self.spectrum_canvas.setFrameStyle(QFrame.Box | QFrame.Raised)
            self.spectrum_canvas.setAlignment(Qt.AlignCenter)
            self.spectrum_canvas.setText("频谱图区域")
            spectrum_layout.addWidget(self.spectrum_canvas, alignment=Qt.AlignCenter)
            viz_right_h.addLayout(spectrum_layout)

            upper_layout.addWidget(viz_right, 1)

            # ========== 下层：左(预测结果) 右(合理建议+日志) ==========
            lower_layout = QHBoxLayout()

            # 下左：预测结果（预测类别 + 概率分布）
            result_group = QGroupBox("预测结果")
            result_v = QVBoxLayout(result_group)

            # 预测类别
            row1 = QHBoxLayout()
            row1.addWidget(QLabel("预测类别："))
            self.pred_class_value = QLabel("未开始")
            row1.addWidget(self.pred_class_value)
            row1.addStretch(1)
            result_v.addLayout(row1)

            # 概率分布（改为 QTextEdit）
            self.prob_text = QTextEdit()
            self.prob_text.setReadOnly(True)
            self.prob_text.setFixedWidth(280)
            self.prob_text.setStyleSheet(
                "QTextEdit { font-family: Consolas, 'Courier New', monospace; font-size: 12px; }")
            self.prob_text.setText(
                "正常轴承         : 0.00%\n"
                "内圈故障轴承     : 0.00%\n"
                "外圈故障轴承     : 0.00%\n"
                "滚动体故障轴承   : 0.00%\n"
            )
            result_v.addWidget(self.prob_text)

            lower_layout.addWidget(result_group, 0, Qt.AlignTop)

            # 下右："合理建议" + 日志框（标题与左侧对齐）
            advice_log_wrapper = QWidget()
            advice_log_v = QVBoxLayout(advice_log_wrapper)

            # 与左侧标题对齐的“合理建议”
            self.advice_label = QLabel("合理建议")
            self.advice_label.setStyleSheet("font-weight: bold; font-size: 14px;")
            advice_log_v.addWidget(self.advice_label)

            # 日志框
            self.log_box = QTextEdit()
            self.log_box.setReadOnly(True)
            self.log_box.setFixedHeight(220)
            advice_log_v.addWidget(self.log_box)

            lower_layout.addWidget(advice_log_wrapper, 1)

            # ========== 装配到整体 ==========
            bottom_layout.addLayout(upper_layout, stretch=2)
            bottom_layout.addLayout(lower_layout, stretch=1)

            main_layout.addWidget(bottom_widget, stretch=2)


        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("Training 线程报错：\n", error_msg)  # 控制台显示

    def start_preprocess_clicked(self):
        try:
            has_data = any(getattr(it, "files", []) for it in self.label_items)
            if not has_data:
                self.output_box.append("请先为各标签加载数据目录，然后再开始预处理。")
                return
            self.output_box.append("开始预处理数据...")
            self.preprocessed_area.start_preprocess()
            self.output_box.append("预处理结束，可以开始训练模型")
        except Exception:
            import traceback
            self.output_box.append("预处理启动失败，请查看控制台。")
            print("start_preprocess_clicked 报错：\n", traceback.format_exc())

    def center_text_edit(self, text_edit: QTextEdit, text: str):
        cursor = text_edit.textCursor()
        fmt = cursor.blockFormat()
        fmt.setAlignment(Qt.AlignCenter)
        cursor.setBlockFormat(fmt)
        text_edit.setTextCursor(cursor)
        text_edit.setText(text)

    '''###### Train ######'''
    def show_selected_models(self):
        try:
            selected = self.model_select.checked_texts()
            self.output_box.append("选择的模型: " + (", ".join(selected) if selected else "（未选择）"))
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("show_selected_models 线程报错：\n", error_msg)  # 控制台显示

    def get_selected_models(self):
        try:
            return self.model_select.checked_texts()
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("get_selected_models 线程报错：\n", error_msg)  # 控制台显示

    def update_labels(self):
        try:
            # 清空现有标签
            for item in self.label_items:
                item.setParent(None)
            self.label_items = []

            n = int(self.label_count_input.text())
            for i in range(1, n + 1):
                item = LabelItem(i)
                self.scroll_layout.addWidget(item)
                self.label_items.append(item)
            self.scroll_layout.addStretch(1)
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("start_training 线程报错：\n", error_msg)  # 控制台显示

    def start_training(self):
        try:
            # 1) 检查是否已经有预处理结果
            if not getattr(self, "preprocessed_data", None):
                self.output_box.append("尚未完成数据预处理，请先点击“开始预处理”。")
                return

            # 2) 读取所选模型
            self.selected_models = self.model_select.checked_texts()
            if not self.selected_models:
                self.output_box.append("请先选择模型")
                return

            # 3) 读取训练超参
            self.ui_epochs = int(self.epoch_spin.value())
            self.ui_lr = float(self.lr_spin.value())
            self.ui_wd = float(self.wd_spin.value())
            self.ui_gamma = float(self.gamma_spin.value())

            self.output_box.append(
                f"初始化模型... epochs={self.ui_epochs}, lr={self.ui_lr}, wd={self.ui_wd}, gamma={self.ui_gamma}"
            )

            # 4) 初始化模型线程
            self.model_thread = ModelInitThread(self.selected_models, self.ui_lr, self.ui_wd, self.ui_gamma)
            self.model_thread.finished.connect(self.on_models_ready)
            self.model_thread.start()

        except Exception:
            import traceback
            print("start_training 线程报错：\n", traceback.format_exc())

    def on_models_ready(self, modell, loss_fn):
        try:
            print(1)
            self.output_box.append("模型初始化完成，可以开始训练")
            self.modell = modell
            self.loss_fn = loss_fn
            self.model_name = list(self.modell.keys())
            print("初始化的模型:", modell.keys())
            self.train_data = self.preprocessed_data  # 保存到类属性
            print("选择的模型:", list(self.modell.keys()))
            print("模型的类别:", list(self.preprocessed_data.keys()))

            # --- 在这里启动数据准备线程 ---
            self.data_thread = DataPrepareThread(self.preprocessed_data)
            self.data_thread.finished.connect(self.on_datasets_ready)
            self.data_thread.start()
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("on_models_ready 报错：\n", error_msg)  # 控制台显示

    def on_datasets_ready(self, GADF_loader, CWT_loader, EMTF_loader):
        try:
            print("数据集准备完成，可以开始训练")
            self.GADF_loader = GADF_loader
            self.CWT_loader = CWT_loader
            self.EMTF_loader = EMTF_loader
            print(self.GADF_loader)
            self.output_box.append("数据集准备完成，可以开始训练")
            print("开始训练")
            # === 确保这里已经有 DataLoader 了 ===
            self.loaders = {
                "GADF": self.GADF_loader,
                "CWT": self.CWT_loader,
                "EMTF": self.EMTF_loader
            }

            if hasattr(self, 'modell'):
                self.eval_panel.update_context(self.modell, self.loaders, self.tsne_max_model_combo)

            # 启动训练线程
            self.train_thread = TrainThread(self.selected_models,
                                            self.modell, self.loss_fn, self.loaders,
                                            epochs=self.ui_epochs)
            self.train_thread.log_signal.connect(self.output_box.append)
            self.train_thread.finished.connect(self.on_training_finished)
            self.train_thread.current_progress.connect(self.text_current_progress.setText)
            self.train_thread.learning_rate.connect(self.text_current_learning_rate.setText)
            self.train_thread.extimated_remaining_time.connect(self.text_estimated_remaining_time.setText)
            self.train_thread.acc_loss.connect(self.acc_loss_pix)

            self.train_thread.start()
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("on_datasets_ready 报错：\n", error_msg)  # 控制台显示
    def on_training_finished(self, results):
        try:
            self.output_box.append("所有训练完成！")
            # 遍历训练完成的模型，把名字加到 T-SNE 下拉框
            for model_name in results.keys():
                if self.model_combo.findText(model_name) == -1:  # 避免重复添加
                    if model_name not in self.tsne_max_model_combo:
                        self.tsne_max_model_combo.append(model_name)

                    self.model_combo.addItem(model_name)

            self.eval_panel.update_context(self.modell, self.loaders, self.tsne_max_model_combo)

            for model, acc in results.items():
                self.output_box.append(f"{model}: 最终准确率 {acc:.4f}")
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("on_training_finished 报错：\n", error_msg)  # 控制台显示

    def set_preprocessed_data(self, results):
        # print("TrainArea: 收到数据", results)
        try:
            self.preprocessed_data = results
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("set_preprocessed_data 报错：\n", error_msg)  # 控制台显示

    def update_log(self, msg: str):
        """把日志信息显示到界面输出框"""
        try:
            self.output_box.append(msg)  # 假设你有一个 QTextEdit 或 QTextBrowser 名为 self.output_box
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("update_log 报错：\n", error_msg)  # 控制台显示

    def show_image_on_label(self, label: QLabel, img_path: str):
        pix = QPixmap(img_path)
        if pix.isNull():
            label.setText("图片加载失败")
            return
        scaled = pix.scaled(label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
        label.setPixmap(scaled)

    def acc_loss_pix(self, payload: bytes):
        pix = QPixmap()
        pix.loadFromData(payload)
        scaled = pix.scaled(self.SQ_area.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
        self.SQ_area.setPixmap(scaled)


    '''############################ 预测区域 ############################'''
    # ===== 文件选择 =====
    def select_file(self):
        path, _ = QFileDialog.getOpenFileName(self, "选择文件")
        if path:
            self.filepath = path
            self.log_box.append(f"选择文件: {path}")

    # ===== 开始预测 =====
    def start_predict(self):
        try:
            model_name = self.model_combo.currentText()
            if model_name == "请选择模型" or not hasattr(self, "filepath"):
                self.log_box.append("请先选择文件和模型！")
                return

            self.log_box.append(f"开始预测，模型={model_name}, 文件={self.filepath}")
            model = self.modell[model_name]["model"]
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

            # 从UI读取参数
            fs = int(self.sr_spin.value())
            seg_len = int(self.seg_spin.value())

            self.predict_thread = PredictThread(model_name, model, self.filepath, device,
                                                fs=fs, seg_len=seg_len)
            self.predict_thread.log_signal.connect(lambda msg: self.log_box.append(msg))
            self.predict_thread.waveform_ready.connect(self._show_waveform_png)
            self.predict_thread.spectrum_ready.connect(self._show_spectrum_png)
            self.predict_thread.final_advice.connect(self._show_final_advice)

            # ▶ 用 result_signal 更新概率分布与预测类别
            self.predict_thread.result_signal.connect(self._update_pred_panel)

            self.predict_thread.start()
        except Exception as e:
            import traceback
            error_msg = traceback.format_exc()
            print("start_predict 线程报错：\n", error_msg)


    def _show_waveform_png(self, payload: bytes):
        try:
            pix = QPixmap()
            pix.loadFromData(payload)
            scaled = pix.scaled(self.waveform_canvas.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.waveform_canvas.setPixmap(scaled)
        except Exception as e:
            import traceback
            print("显示时域波形失败：\n", traceback.format_exc())

    def _show_spectrum_png(self, payload: bytes):
        try:
            pix = QPixmap()
            pix.loadFromData(payload)
            scaled = pix.scaled(self.spectrum_canvas.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.spectrum_canvas.setPixmap(scaled)
        except Exception as e:
            import traceback
            print("显示频谱图失败：\n", traceback.format_exc())

    def _show_final_advice(self, label_name: str, advice: str):
        # 清晰一点的输出格式
        self.log_box.append(f"\n【最终判定】{label_name}")
        self.log_box.append(advice)

    def _update_pred_panel(self, preds_list):
        from collections import Counter
        id2name = {0: "正常轴承", 1: "内圈故障轴承", 2: "外圈故障轴承", 3: "滚动体故障轴承"}

        if not preds_list:
            self.pred_class_value.setText("无法判定")
            self.prob_text.setPlainText(
                "正常轴承         : 0.00%\n"
                "内圈故障轴承     : 0.00%\n"
                "外圈故障轴承     : 0.00%\n"
                "滚动体故障轴承   : 0.00%\n"
            )
            return

        # 再保守一次：把任何 numpy 类型都转成 Python int
        preds_list = [int(x) for x in preds_list]

        cnt = Counter(preds_list)
        total = float(sum(cnt.values()))
        final_id = cnt.most_common(1)[0][0]
        self.pred_class_value.setText(id2name.get(final_id, f"未知({final_id})"))

        lines = []
        for cid in (0, 1, 2, 3):
            p = 100.0 * cnt.get(cid, 0) / total
            lines.append(f"{id2name[cid]:<12} : {p:5.2f}%")
        self.prob_text.setPlainText("\n".join(lines))


class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("训练与预测界面")
        self.resize(1500, 1500)

        # 主布局
        main_layout = QVBoxLayout(self)
        # 左侧训练部分
        train_frame = QFrame()
        train_layout = QVBoxLayout(train_frame)
        self.dataset_prepare = TrainArea()
        train_layout.addWidget(self.dataset_prepare, stretch=0)

        # 添加到主布局
        main_layout.addWidget(train_frame, 3)
        self.setLayout(main_layout)


if __name__ == "__main__":
    import torch
    # -------------------- 提前初始化 CUDA --------------------
    if torch.cuda.is_available():
        print("CUDA 初始化中...")
        torch.empty(1).to("cuda")  # 触发一次 CUDA runtime 初始化
        print("CUDA 初始化完成")
    # ---------------------------------------------------------
    app = QApplication(sys.argv)
    w = MainWindow()
    w.show()
    sys.exit(app.exec_())

